#!/usr/bin/python3
# -*- coding: utf-8 -*-

"""
# @version: v1.0
# @author : cd
# @Email : 19688513@qq.com
# @Project : horizons-engine-pybroker
# @File : StockBasicInfoManager.py
# @Software: PyCharm
# @time: 2025/6/11 15:38
# @description : 股票基本信息管理类（优化版）
"""
import psutil
import urllib3
import logging
from datetime import datetime, time
import pandas as pd
import os
import pickle
from tqdm import tqdm
from diskcache import Cache
from typing import Dict, Optional, List
import re
import sqlite3
import zlib
import msgpack
import numpy as np
from concurrent.futures import ThreadPoolExecutor, as_completed

from AKShareDataFetcher import AKShareDataFetcher

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
# 禁用不安全的请求警告
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


class StockBasicInfoManager:
    """股票基本信息管理类（高性能优化版）"""

    def __init__(self, cache_directory: str = 'stock_basic_cache',
                 db_file: str = 'data/stock_basic_data.db',
                 max_workers: int = 8,
                 auto_initialize: bool = True):  # 新增自动初始化参数
        self.fetcher = AKShareDataFetcher()  # 初始化数据获取器
        self.cache = Cache(cache_directory)  # 创建Diskcache实例
        self.db_file = db_file
        self.max_workers = max_workers

        # 确保数据库目录存在
        os.makedirs(os.path.dirname(db_file), exist_ok=True)

        # 初始化数据库连接
        self.conn = self._create_connection()
        self._initialize_database()

        # 内存缓存
        self.stock_cache = {}
        self._load_stock_cache()

        # 自动初始化数据库（新增）
        if auto_initialize and not self.stock_cache:
            logging.warning("数据库为空，自动初始化股票数据...")
            self.initialize_stock_database()
            self._load_stock_cache()  # 重新加载缓存

    def _create_connection(self):
        """创建高性能数据库连接"""
        try:
            # 确保数据库文件目录存在
            db_dir = os.path.dirname(self.db_file)
            logging.info(f"数据库目录: {db_dir}")

            if db_dir and not os.path.exists(db_dir):
                logging.info(f"创建数据库目录: {db_dir}")
                os.makedirs(db_dir, exist_ok=True)

            # 检查目录是否存在
            if db_dir and os.path.exists(db_dir):
                logging.info(f"数据库目录已存在: {db_dir}")
            else:
                logging.warning("数据库目录为空，将在当前目录创建数据库")

            # 检查数据库文件是否存在
            db_exists = os.path.exists(self.db_file)
            logging.info(f"数据库文件 {self.db_file} 存在: {db_exists}")

            # 创建数据库连接
            conn = sqlite3.connect(self.db_file)
            conn.execute("PRAGMA journal_mode = WAL;")
            conn.execute("PRAGMA synchronous = NORMAL;")
            conn.execute("PRAGMA cache_size = -20000;")  # 20MB缓存
            conn.execute("PRAGMA temp_store = MEMORY;")

            # 检查连接是否成功
            try:
                cursor = conn.cursor()
                cursor.execute("SELECT 1")
                cursor.fetchone()
                logging.info("数据库连接成功")
            except Exception as test_e:
                logging.error(f"数据库连接测试失败: {str(test_e)}")
                # 如果测试失败，尝试重新连接
                conn.close()
                conn = sqlite3.connect(self.db_file)
                logging.info("已尝试重新连接数据库")

            return conn
        except Exception as e:
            logging.error(f"创建数据库连接失败: {str(e)}", exc_info=True)
            # 创建内存数据库作为后备
            logging.warning("将使用内存数据库作为后备方案")
            return sqlite3.connect(":memory:")

    def _initialize_database(self):
        """初始化数据库表结构"""
        try:
            # 检查表是否已存在
            cursor = self.conn.cursor()
            cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='stock_basic_info'")
            table_exists = cursor.fetchone() is not None

            if not table_exists:
                logging.info("创建 stock_basic_info 表...")
                self.conn.execute('''
                    CREATE TABLE stock_basic_info (
                        stock_code TEXT PRIMARY KEY,
                        stock_name TEXT NOT NULL,
                        total_market_value REAL,
                        circulating_market_value REAL,
                        industry TEXT,
                        listing_date TEXT,
                        total_shares INTEGER,
                        circulating_shares INTEGER,
                        rating_agency TEXT,
                        rating REAL,
                        rating_quarter TEXT,
                        market_type TEXT,
                        last_updated TEXT
                    )
                ''')

            # 检查索引是否已存在
            cursor.execute("PRAGMA index_list('stock_basic_info')")
            indexes = [row[1] for row in cursor.fetchall()]

            if 'idx_stock_code' not in indexes:
                logging.info("创建 idx_stock_code 索引...")
                self.conn.execute('CREATE INDEX idx_stock_code ON stock_basic_info(stock_code)')

            if 'idx_market_type' not in indexes:
                logging.info("创建 idx_market_type 索引...")
                self.conn.execute('CREATE INDEX idx_market_type ON stock_basic_info(market_type)')

            self.conn.commit()
            logging.info("数据库表结构初始化成功")
            return True
        except Exception as e:
            logging.error(f"初始化数据库表结构失败: {str(e)}", exc_info=True)
            self.conn.rollback()
            return False

    def _load_stock_cache(self):
        """从数据库加载股票数据到内存缓存（修复字段缺失问题）"""
        try:
            # 检查表是否存在
            cursor = self.conn.cursor()
            cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='stock_basic_info'")
            table_exists = cursor.fetchone() is not None

            if not table_exists:
                logging.warning("stock_basic_info表不存在，跳过加载")
                self.stock_cache = {}
                return False

            df = pd.read_sql_query("SELECT * FROM stock_basic_info", self.conn)
            if not df.empty:
                # 转换为字典 {stock_code: row_dict}
                self.stock_cache = df.set_index('stock_code').to_dict('index')

                # 修复：确保所有条目都有stock_code字段
                for code, info in self.stock_cache.items():
                    info['stock_code'] = code

                logging.info(f"从数据库加载了 {len(self.stock_cache)} 只股票数据到缓存")
                return True
            else:
                logging.info("数据库为空，将从头开始初始化")
                self.stock_cache = {}
                return False
        except Exception as e:
            logging.error(f"加载数据库失败: {str(e)}", exc_info=True)
            self.stock_cache = {}
            return False

    def compress_data(self, data):
        """压缩数据"""
        return zlib.compress(msgpack.packb(data), zlib.Z_BEST_COMPRESSION)

    def decompress_data(self, compressed_data):
        """解压缩数据"""
        return msgpack.unpackb(zlib.decompress(compressed_data))

    def _convert_to_int(self, value: str) -> int:
        """将字符串值转换为整数（处理带单位的数值）"""
        try:
            if isinstance(value, (int, float)):
                return int(value)

            # 处理带单位的数值
            if '亿' in value:
                num = float(value.replace('亿', '').replace(',', ''))
                return int(num * 100000000)
            elif '万' in value:
                num = float(value.replace('万', '').replace(',', ''))
                return int(num * 10000)
            return int(float(value.replace(',', '')))
        except Exception:
            return 0

    def _convert_to_float(self, value: str) -> float:
        """将字符串值转换为浮点数（处理带单位的数值）"""
        try:
            if isinstance(value, (int, float)):
                return float(value)

            # 处理带单位的数值
            if '亿' in value:
                num = float(value.replace('亿', '').replace(',', ''))
                return num * 100000000
            elif '万' in value:
                num = float(value.replace('万', '').replace(',', ''))
                return num * 10000
            return float(value.replace(',', ''))
        except Exception:
            return 0.0

    def get_stock_exchange(self, code: str) -> str:
        """根据股票代码判断所属交易所"""
        code_str = str(code).zfill(6)
        if code_str.startswith(('6', '9')):
            return 'sh'
        elif code_str.startswith(('0', '3')):
            return 'sz'
        elif code_str.startswith('4') or code_str.startswith('8'):
            return 'bj'
        else:
            return 'unknown'

    def get_stock_info_by_code(self, stock_code: str) -> Optional[dict]:
        """根据股票代码获取股票信息（优先从缓存获取）"""
        # 确保股票代码是字符串格式
        stock_code = str(stock_code).strip()
        return self.stock_cache.get(stock_code)

    def get_stock_codes(self, market_types: list = None, excluded_prefixes: list = None) -> pd.DataFrame:
        """获取符合条件的股票代码（添加健壮性检查）"""
        # 设置默认值
        if market_types is None:
            market_types = ['sz', 'sh']
        if excluded_prefixes is None:
            excluded_prefixes = ['688', '300', '83', '87', '43']  # 含科创板、创业板、北交所


        # 从缓存中筛选数据
        filtered_data = []
        for code, data in self.stock_cache.items():
            # 确保数据包含必要字段
            if 'market_type' not in data:
                continue

            # 只选择符合市场类型和排除前缀的股票
            if data['market_type'] not in market_types:
                continue
            if any(str(code).startswith(prefix) for prefix in excluded_prefixes):
                continue

            # 创建包含stock_code的新字典
            item = {
                'stock_code': code,  # 显式添加stock_code字段
                'stock_name': data.get('stock_name', ''),
                'total_market_value': data.get('total_market_value', 0.0),
                'circulating_market_value': data.get('circulating_market_value', 0.0),
                'industry': data.get('industry', ''),
                'listing_date': data.get('listing_date', ''),
                'total_shares': data.get('total_shares', 0),
                'circulating_shares': data.get('circulating_shares', 0),
                'rating_agency': data.get('rating_agency', ''),
                'rating': data.get('rating', 0.0),
                'rating_quarter': data.get('rating_quarter', ''),
                'market_type': data.get('market_type', 'unknown'),
                'last_updated': data.get('last_updated', '')
            }
            filtered_data.append(item)

        # 转换为DataFrame并设置列顺序
        if filtered_data:
            df = pd.DataFrame(filtered_data)
            cols = ['stock_code', 'listing_date', 'stock_name', 'industry',
                    'total_shares', 'circulating_shares', 'total_market_value',
                    'circulating_market_value', 'market_type', 'rating_agency',
                    'rating', 'rating_quarter', 'last_updated']
            return df[cols]  # 保持列顺序一致性
        else:
            # 返回空的DataFrame
            columns = ['stock_code', 'listing_date', 'stock_name', 'industry',
                       'total_shares', 'circulating_shares', 'total_market_value',
                       'circulating_market_value', 'market_type', 'rating_agency',
                       'rating', 'rating_quarter', 'last_updated']
            return pd.DataFrame(columns=columns)

    def _clean_stock_code(self, raw_code: str) -> str:
        """清洗股票代码，专注于A股市场"""
        try:
            # 统一代码清洗
            code_str = str(raw_code).strip().upper()

            # 移除所有市场后缀（如.SH, .SZ, .HK等）
            cleaned = re.sub(r'\.[A-Z]+$', '', code_str)

            # 只处理A股股票（6位纯数字代码）
            if cleaned.isdigit() and len(cleaned) == 6:
                return cleaned

            # 非A股股票，返回None跳过
            return None
        except Exception as e:
            logging.debug(f"代码清洗失败: {str(e)}", exc_info=False)
            return None

    def _build_esg_mapping(self, esg_df):
        """构建ESG映射字典（优化版）"""
        mapping = {}
        if esg_df is not None and not esg_df.empty:
            # 构建映射（优先使用最新评级）
            for _, row in esg_df.iterrows():
                raw_code = row.get('股票代码', '')
                code = self._clean_stock_code(raw_code)
                if not code:
                    continue

                # 只保留最新评级
                if code not in mapping:
                    mapping[code] = {
                        'ESG评分': row.get('ESG评分', "无"),
                        '交易市场': row.get('交易市场', "无"),
                        'ESG等级': row.get('ESG等级', '未评级')
                    }
        return mapping

    def _fetch_stock_detail(self, stock_code: str, stock_name: str, esg_mapping: dict) -> bool:
        """获取并处理单只股票详细信息"""
        try:
            # 尝试从缓存中获取股票详细信息
            cache_key = f'stock_detail_info_{stock_code}'
            info_df = self.cache.get(cache_key)

            if info_df is None or info_df.empty:
                # 如果缓存中没有数据，则从数据源获取
                info_df = self.fetcher.fetch_stock_detail_info(stock_code)
                if not info_df.empty:
                    # 缓存数据，有效期24小时
                    self.cache.set(cache_key, info_df, expire=3600 * 24)
                else:
                    logging.warning(f"无法获取股票 {stock_code} 的详细信息")
                    return False

            # 转换为字典格式
            info_dict = info_df.set_index('item')['value'].to_dict()

            # 处理ESG信息
            esg_info = esg_mapping.get(stock_code, {})
            rating_agency = "华证" if esg_info else "无"
            rating = esg_info.get('ESG评分', "无")
            rating_quarter = esg_info.get('ESG等级', '未评级')

            # 构造股票信息（添加默认值）
            stock_info = {
                'stock_code': stock_code,  # 确保包含股票代码
                'stock_name': stock_name,
                'total_market_value': self._convert_to_float(info_dict.get('总市值', 0)),
                'circulating_market_value': self._convert_to_float(info_dict.get('流通市值', 0)),
                'industry': info_dict.get('行业', '未知'),
                'listing_date': info_dict.get('上市时间', ''),
                'total_shares': self._convert_to_int(info_dict.get('总股本', '0')),
                'circulating_shares': self._convert_to_int(info_dict.get('流通股', '0')),
                'rating_agency': rating_agency,
                'rating': rating,
                'rating_quarter': rating_quarter,
                'market_type': self.get_stock_exchange(stock_code),
                'last_updated': datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            }

            # 更新内存缓存
            self.stock_cache[stock_code] = stock_info
            return True
        except Exception as e:
            logging.error(f"股票 {stock_code} 基本信息更新失败: {str(e)}", exc_info=True)
            return False

    def _save_to_database(self):
        """批量保存股票数据到数据库（修复KeyError问题）"""
        if not self.stock_cache:
            return False

        try:
            # 准备数据：使用字典键作为股票代码
            data = []
            for stock_code, stock_info in self.stock_cache.items():
                # 添加健壮性检查，确保必要字段存在
                data.append((
                    stock_code,  # 使用字典键作为股票代码
                    stock_info.get('stock_name', ''),
                    stock_info.get('total_market_value', 0.0),
                    stock_info.get('circulating_market_value', 0.0),
                    stock_info.get('industry', ''),
                    stock_info.get('listing_date', ''),
                    stock_info.get('total_shares', 0),
                    stock_info.get('circulating_shares', 0),
                    stock_info.get('rating_agency', ''),
                    stock_info.get('rating', 0.0),
                    stock_info.get('rating_quarter', ''),
                    stock_info.get('market_type', 'unknown'),
                    stock_info.get('last_updated', datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
                ))

            # 批量插入或更新
            self.conn.executemany('''
                INSERT OR REPLACE INTO stock_basic_info 
                (stock_code, stock_name, total_market_value, circulating_market_value, 
                 industry, listing_date, total_shares, circulating_shares, 
                 rating_agency, rating, rating_quarter, market_type, last_updated)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            ''', data)
            self.conn.commit()
            logging.info(f"成功保存 {len(data)} 条股票数据到数据库")
            return True
        except Exception as e:
            logging.error(f"保存到数据库失败: {str(e)}", exc_info=True)
            self.conn.rollback()
            return False

    def initialize_stock_database(self, force_refresh: bool = False):
        """高性能初始化股票数据库"""
        logging.info("🏁 开始高性能初始化股票数据库...")
        start_time = datetime.now()
        success_count = 0
        failed_stocks = []
        total_count = 0  # 初始化为0

        try:
            # 获取股票基本信息
            basic_info = self.fetcher.fetch_stock_basic_info(
                max_retries=3,
                retry_delay=2,
                request_delay=0.2
            )
            if basic_info is None or basic_info.empty:
                logging.error("股票基本信息获取失败")
                return

            # 获取ESG信息
            esg_info_df = self.fetcher.fetch_stock_esg_info(
                max_retries=3,
                retry_delay=3,
                request_delay=0.3
            )
            esg_mapping = self._build_esg_mapping(esg_info_df)

            # 准备股票列表
            stock_list = []
            for _, row in basic_info.iterrows():
                stock_code = self._clean_stock_code(row['code'])
                if not stock_code:
                    continue
                stock_list.append((stock_code, row['name']))

            total_count = len(stock_list)
            logging.info(f"开始处理 {total_count} 只股票...")

            # 使用线程池并行处理
            with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                futures = []
                for stock_code, stock_name in stock_list:
                    futures.append(
                        executor.submit(
                            self._fetch_stock_detail,
                            stock_code,
                            stock_name,
                            esg_mapping
                        )
                    )

                # 处理结果
                for future in tqdm(as_completed(futures), total=total_count, desc="🚀 并行处理股票"):
                    try:
                        if future.result():
                            success_count += 1
                        else:
                            # 记录失败的股票代码
                            failed_stocks.append(stock_code)
                    except Exception as e:
                        failed_stocks.append(str(e))

            # 保存到数据库
            if not self._save_to_database():
                logging.error("保存股票数据到数据库失败")

        except Exception as e:
            logging.error(f"初始化失败: {str(e)}", exc_info=True)
            return

        finally:
            # 性能统计
            elapsed = datetime.now() - start_time
            mem_usage = psutil.Process().memory_info().rss / 1024 ** 2

            # 结构化日志
            log_data = {
                '耗时': str(elapsed),
                '总股票数': total_count,
                '成功': success_count,
                '失败': len(failed_stocks),
                '内存占用': f"{mem_usage:.2f}MB",
                '线程数': self.max_workers
            }
            logging.info("📊 初始化完成统计: %s", log_data)

            if failed_stocks:
                logging.warning(
                    "⚠️ 部分股票更新失败: %s",
                    ', '.join(failed_stocks[:20]) + ('...' if len(failed_stocks) > 20 else '')
                )

    def get_stock_count(self) -> int:
        """获取当前存储的股票数量"""
        return len(self.stock_cache)

    def clear_cache(self):
        """清除缓存"""
        self.cache.clear()
        logging.info("磁盘缓存已清除")

    def __del__(self):
        """析构函数，确保数据库连接关闭"""
        try:
            if hasattr(self, 'conn') and self.conn:
                self.conn.close()
                logging.info("数据库连接已关闭")
        except Exception as e:
            logging.error(f"关闭数据库连接失败: {str(e)}")


if __name__ == "__main__":
    # 创建股票基本信息管理实例
    stock_manager = StockBasicInfoManager(
        max_workers=12
    )

    # 检查数据库是否初始化成功
    if stock_manager._initialize_database():
        logging.info("数据库已正确初始化")
    else:
        logging.error("数据库初始化失败")

    # 初始化股票数据库
    # stock_manager.initialize_stock_database()

    # 示例：获取符合条件的股票代码
    logging.info("获取符合条件的股票代码...")
    filtered_stock_codes = stock_manager.get_stock_codes()

    logging.info(f"符合条件的股票数量: {len(filtered_stock_codes)}")
    logging.info(f"当前存储的股票总数: {stock_manager.get_stock_count()}")

    # 打印前5只股票
    if not filtered_stock_codes.empty:
        print(filtered_stock_codes.head(5))

    # 示例：获取特定股票信息
    test_stock = "000001"
    stock_info = stock_manager.get_stock_info_by_code(test_stock)
    if stock_info:
        logging.info(f"股票 {test_stock} 的信息:")
        for key, value in stock_info.items():
            logging.info(f"{key}: {value}")
    else:
        logging.warning(f"未找到股票 {test_stock} 的信息")